#!/usr/bin/python
'''
    Filename - ftps.py
    Author: Gerrit Kitts <gkitts@gmail.com>        
    Usage: ftps.py (listenPort)
    
    This package provides a very basic server that will listen on a specified port
    for incoming connections and accept binary file transfers from a remote host via UDP.
    Each UDP packet consists of:
        remote IP <4 bytes>: The destination IP for the datagram
        remote Port <2 bytes>: The destination port for the datagram
        datagram type <1 byte>: 1 if the datagram contains the file size, 2 if the datagram
            contains the file name, 3 for data segments.
'''

import array
import hashlib
import socket
import struct
import os
import os.path

from sys import argv, exit

FLAG_SIZE = 1
FLAG_NAME = 2
FLAG_DATA = 3

SEQ_ZERO = 0
SEQ_ONE = 1

BUFF_SIZE = 1025
DATA_SIZE = 1000
HEADER_STRUCT = '! 4B H B B' # common header elements for packer functions


'''
    Exception sub-class to handle invalid argument specifications
'''
class ArgumentException(BaseException):
    def __init__(self):
        self.msg = "Usage: ftps.py (remoteIP) (remotePort) (file-to-transfer)"

'''
    Extracts the flag value from the datagram header and returns it.
'''
def determineFlagAndSeq(data):
    unpacker = struct.Struct('! 4B H B B')
    metadata = unpacker.unpack(data[0:8])
    destIP = '.'.join([str(x) for x in metadata[0:4]])
    destPort = metadata[4]
    flag = metadata[5]
    seq = metadata[6]
    return (flag, seq)


'''
    Creates a socket on all IP addresses on the port specified as a parameter, and listens
    for incoming requests.  
'''
def listenForConnections(params):
    host = socket.gethostbyname(socket.gethostname())  
    localIP = [int(x) for x in host.split('.')]
    localPort = params['address'][1]
    localAddrS = ".".join([str(x) for x in localIP])
    trollAddr = ('127.0.0.1', params['trollPort'])
    print "Local addr is %s:%s" % (localAddrS, localPort)
    print "Troll addr is %s:%s" % (trollAddr[0], trollAddr[1])
    
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    address = params['address']
    s.bind(address)
    
    # Track our current state for a expected single set of datagrams in flight at 
    #  any given point in time.
    state = 1
    size = -1
    remaining = -1
    fname = ""    
    m = hashlib.md5() # hash for the MD5 checksum
    f = ""

    '''
        Note: This never terminates.  This was stated as being acceptable on the 
        group, so if there are dropped packets during the data read we'll persist
        indefinitely.  However, we do attempt to track the state of the name & size
        reads and fail if those occur 'out of order'
    '''
    expectedSeq = -1
    while True:
        data = s.recv(BUFF_SIZE)

        # First, we have to figure out what the state of this datagram is, so decode it
        (flag, seq) = determineFlagAndSeq(data)
        if seq is not expectedSeq and expectedSeq is not -1:
            # if seq != -1, we've sent an ACK but gotten the same packet again.  Resend the ACK.
            packer = struct.Struct(HEADER_STRUCT)
            remIP = list(metadata[0:4])
            remPort = metadata[4]
            remAddr = "%s:%s" % (".".join([str(x) for x in remIP]), remPort)            
            print " SEQ was %i but expected %i.  Resending ACK resending to %s." % (seq, expectedSeq, remAddr)
            sizeACK = packer.pack(*(remIP + [remPort, FLAG_SIZE, seq]))
            s.sendto(sizeACK, trollAddr)
            continue

        if flag == state:
            # Continue down the primrose path
            if flag == FLAG_SIZE:
                # If we get the size, set the remaining bytes (for tracking) and go
                unpacker = struct.Struct(HEADER_STRUCT + 'I')
                metadata = unpacker.unpack(data)
                size = metadata[7]
                remaining = size
                
                # Send an ACK for the size
                packer = struct.Struct(HEADER_STRUCT)
                remIP = list(metadata[0:4])
                remPort = metadata[4]
                remAddr = "%s:%s" % (".".join([str(x) for x in remIP]), remPort)
                #print " Remote IP:Port is %s:%s" % (remIP, remPort)
                sizeACK = packer.pack(*(remIP + [remPort, FLAG_SIZE, SEQ_ZERO]))
                print " Sending ACK 0 for file size to %s" % (remAddr),
                s.sendto(sizeACK, trollAddr)
                expectedSeq = 1
                
                # Update the state
                state = FLAG_NAME 
                               
            elif flag == FLAG_NAME:
                # If we have the filename, create the output file
                unpacker = struct.Struct(HEADER_STRUCT + '20s')
                metadata = unpacker.unpack(data)
                name = metadata[7].rstrip('\0')
                fname = os.path.join(params['dest'], name)
                print " File will be written as %s" % (fname)
                f = open(fname, 'wb')
                
                # Send an ACK for the filename
                packer = struct.Struct(HEADER_STRUCT)
                remIP = list(metadata[0:4])
                remPort = metadata[4]
                remAddr = "%s:%s" % (".".join([str(x) for x in remIP]), remPort)
                #print " Remote IP:Port is %s:%s" % (remIP, remPort)
                sizeACK = packer.pack(*(remIP + [remPort, FLAG_SIZE, SEQ_ONE]))
                print " Sending ACK 1 for file name to %s" % (remAddr),
                s.sendto(sizeACK, trollAddr)                
                expectedSeq = 0
                                
                # Update the state
                state = FLAG_DATA
                
            elif flag == FLAG_DATA:
                # If we've received data, unpack it and write to the file
                unpacker = struct.Struct(HEADER_STRUCT + '1000s')
                metadata = unpacker.unpack(data)
                seq = metadata[6]
                bytes = metadata[7]
                if remaining < DATA_SIZE:
                    bytes = bytes[0:remaining]
                    remaining = 0
                else:
                    remaining = remaining - DATA_SIZE
                f.write(bytes)
                m.update(bytes)
                
                # Send an ACK for the data
                packer = struct.Struct(HEADER_STRUCT)
                remIP = list(metadata[0:4])
                remPort = metadata[4]
                remAddr = "%s:%s" % (".".join([str(x) for x in remIP]), remPort)                
                #print " Remote IP:Port is %s:%s" % (remIP, remPort)
                sizeACK = packer.pack(*(remIP + [remPort, FLAG_SIZE, seq]))
                print " Sending ACK %i for data chunk to %s" % (seq, remAddr),
                s.sendto(sizeACK, trollAddr)                
                
                # Update the expected sequence
                if seq is 0:
                    expectedSeq = 1
                else:
                    expectedSeq = 0
            
            # Check to see if we've downloaded all the necessary bytes
            if remaining == 0:
                print "\nFile downloaded successfully, ",
                
                # Close the output file handle 
                f.close()

                # Write the MD5checksum
                print "MD5 checksum is: %s" % (m.hexdigest())
                print "Press Ctrl+C to terminate."
            else:
                print " %i bytes remaining" % remaining
            
        else:
            # Something has gone horribly wrong!
            print "Received a packet for flag %i, something has failed.  Aborting!" % (flag)            
            flag = -1
            exit()


'''
    Base wrapper for class functionality
'''
if __name__ == '__main__':
    params = {
        'address' : ('', 7779),
        'trollPort' : 7775,
        'dest' : 'output'
    }

    if len(argv) is 2:
        params['address'] = ('', int(argv[1]))
        
    if len(argv) is 3:
        dirname = argv[2]
        params['dest'] = os.path.abspath(dirname)
        print "Writing files to %s" % (params['dest'])
    
    if len(argv) is 4:
        params['trollPort'] = argv[3]
        print "Using a troll port of: %i" % params['trollPort']
    
    # If the dest path doesn't exist, create it (if possible)
    if not os.path.exists(params['dest']):
        os.mkdir(params['dest'])

    listenForConnections(params)
    
